#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <omp.h>

//Doing just keccak[1600] for now...

#define STATE_ROWS 5

const int rot_const[5][5] = {
	{0, 36, 3, 41, 18},
	{1, 44, 10, 45, 2},
	{62, 6, 43, 15, 61},
	{28, 55, 25, 21, 56},
	{27, 20, 39, 8, 14}
};

inline uint64_t rot(uint64_t in, int n)
{
	return (in << n) | (in >> (64 - n));
}

inline void Round(uint64_t* A, uint64_t RC)
{
	int x,y;
	uint64_t B[STATE_ROWS*STATE_ROWS],C[STATE_ROWS],D[STATE_ROWS];
	
	//θ step
	for(x = 0; x < STATE_ROWS; x++)
	{
		C[x] = A[x*STATE_ROWS] ^ A[x*STATE_ROWS+1] ^ A[x*STATE_ROWS+2] ^ A[x*STATE_ROWS+3] ^ A[x*STATE_ROWS+4];
	}
	for(x = 0; x < STATE_ROWS; x++)
	{
		D[x] = C[(x+4)%STATE_ROWS] ^ rot(C[(x+1)%STATE_ROWS],1);
	}
	for(x = 0; x < STATE_ROWS; x++)
	{
		for(y = 0; y < STATE_ROWS; y++)
		{
			A[x*STATE_ROWS+y] = A[x*STATE_ROWS+y] ^ D[x];
		}
	}

	//ρ and π steps
	for(x = 0; x < 5; x++)
	{
		for(y = 0; y < 5; y++)
		{
			B[y*STATE_ROWS+((2*x+3*y)%STATE_ROWS)] = rot(A[x*STATE_ROWS+y], rot_const[x][y]);
		}
	}

	//χ step
	for(x = 0; x < 5; x++)
	{
		for(y = 0; y < 5; y++)
		{
			A[x*STATE_ROWS+y] = B[x*STATE_ROWS+y] ^ ((~B[((x+1)%STATE_ROWS)*STATE_ROWS+y]) & B[((x+2)%STATE_ROWS)*STATE_ROWS+y]);
		}
	}

	//ι step
	A[0] = A[0] ^ RC;
	
	//return A;
}

void keccak_func(uint64_t *A)
{
	int i;
	uint64_t RC[24] = { 0x0000000000000001,0x0000000000008082,0x800000000000808A,0x8000000080008000,0x000000000000808B,0x0000000080000001,
						0x8000000080008081,0x8000000000008009,0x000000000000008A,0x0000000000000088,0x0000000080008009,0x000000008000000A,
						0x000000008000808B,0x800000000000008B,0x8000000000008089,0x8000000000008003,0x8000000000008002,0x8000000000000080,
						0x000000000000800A,0x800000008000000A,0x8000000080008081,0x8000000000008080,0x0000000080000001,0x8000000080008008
					};
	
	for(i = 0; i < 24; i++)
	{
		int x,y;
	uint64_t B[STATE_ROWS*STATE_ROWS],C[STATE_ROWS],D[STATE_ROWS];
	
	//θ step
	for(x = 0; x < STATE_ROWS; x++)
	{
		C[x] = A[x*STATE_ROWS] ^ A[x*STATE_ROWS+1] ^ A[x*STATE_ROWS+2] ^ A[x*STATE_ROWS+3] ^ A[x*STATE_ROWS+4];
	}
	for(x = 0; x < STATE_ROWS; x++)
	{
		D[x] = C[(x+4)%STATE_ROWS] ^ rot(C[(x+1)%STATE_ROWS],1);
	}
	for(x = 0; x < STATE_ROWS; x++)
	{
		for(y = 0; y < STATE_ROWS; y++)
		{
			A[x*STATE_ROWS+y] = A[x*STATE_ROWS+y] ^ D[x];
		}
	}

	//ρ and π steps
	for(x = 0; x < 5; x++)
	{
		for(y = 0; y < 5; y++)
		{
			B[y*STATE_ROWS+((2*x+3*y)%STATE_ROWS)] = rot(A[x*STATE_ROWS+y], rot_const[x][y]);
		}
	}

	//χ step
	for(x = 0; x < 5; x++)
	{
		for(y = 0; y < 5; y++)
		{
			A[x*STATE_ROWS+y] = B[x*STATE_ROWS+y] ^ ((~B[((x+1)%STATE_ROWS)*STATE_ROWS+y]) & B[((x+2)%STATE_ROWS)*STATE_ROWS+y]);
		}
	}

	//ι step
	A[0] = A[0] ^ RC[i];
	}
}

inline uint64_t reverse(uint64_t in)
{
	return (in << 56) | ((in & 0xFF00) << 40) | ((in & 0xFF0000) << 24) | ((in & 0xFF000000) << 8) | ((in & 0xFF00000000) >> 8) | ((in & 0xFF0000000000) >> 24) | ((in & 0xFF000000000000) >> 40) | ((in & 0xFF00000000000000) >> 56);
}

int main(int argc, char* argv[])
{
	uint8_t *init_message;
	uint8_t *padded_message;
	uint64_t* blocks;
	uint64_t message_length, padded_length;
	FILE *infile;
	struct stat buf;
	int i, bytes_to_pad, x, y, j;
	int count = 0;
	char output[25*16+1] = "";
	char temp[17];
	uint64_t B[STATE_ROWS*STATE_ROWS],C[STATE_ROWS],D[STATE_ROWS];
	uint64_t S[STATE_ROWS*STATE_ROWS] = {0};
	uint64_t RC[24] = { 0x0000000000000001,0x0000000000008082,0x800000000000808A,0x8000000080008000,0x000000000000808B,0x0000000080000001,
						0x8000000080008081,0x8000000000008009,0x000000000000008A,0x0000000000000088,0x0000000080008009,0x000000008000000A,
						0x000000008000808B,0x800000000000008B,0x8000000000008089,0x8000000000008003,0x8000000000008002,0x8000000000000080,
						0x000000000000800A,0x800000008000000A,0x8000000080008081,0x8000000000008080,0x0000000080000001,0x8000000080008008
					};
	double temp_time, profile_time = 0;
	double start;
	if(argc != 2)
	{
		printf("Usage: ./keccak <filename>\n");
		exit(1);
	}
	
	infile = fopen(argv[1], "rb");
	
	if(infile == NULL)
	{
		printf("Could not open input file %s.\n",argv[1]);
		exit(1);
	}
			
	fstat(fileno(infile), &buf);
	
	message_length = buf.st_size;
	//message_length = 1;

	//init_message = (uint8_t*)malloc(message_length*sizeof(uint8_t));
	bytes_to_pad = 72 - (message_length % 72);
	padded_length = message_length + bytes_to_pad;
	padded_message = (uint8_t*)calloc(padded_length,sizeof(uint8_t));
	
	fread((void*)(padded_message),1,message_length,infile);
	
	//printf("Message length = %d\nBytes to add = %d\nPadded length = %d\n",message_length,bytes_to_pad,padded_length);
	blocks = (uint64_t*)malloc(padded_length);
	
	//memcpy(padded_message,init_message,message_length);
	
	padded_message[message_length] = 1;
	padded_message[padded_length-1] |= 0x80;
	memcpy(blocks,padded_message,padded_length);
	/*for(i = 0; i < padded_length; i+=8)
	{
		blocks[i/8] = (uint64_t)padded_message[i] | ((uint64_t)padded_message[i+1] << 8) | ((uint64_t)padded_message[i+2] << 16) | ((uint64_t)padded_message[i+3] << 24) | ((uint64_t)padded_message[i+4] << 32) | ((uint64_t)padded_message[i+5] << 40) | ((uint64_t)padded_message[i+6] << 48) | ((uint64_t)padded_message[i+7] << 56);
	}
	for(i = 0; i < 9; i++)
	{
		printf("blocks: %.16lx\n",blocks[i]);
	}*/
	/*for(i = 0; i < padded_length; i++)
	{
		printf("%.2x",padded_message[i]);
	}
	printf("\n\n");*/
	/*for(i = 0; i < padded_length; i+=8)
	{
		//blocks[i/8] = (uint64_t)padded_message[i] | ((uint64_t)padded_message[i+1] << 8) | ((uint64_t)padded_message[i+2] << 16) | ((uint64_t)padded_message[i+3] << 24) | ((uint64_t)padded_message[i+4] << 32) | ((uint64_t)padded_message[i+5] << 40) | ((uint64_t)padded_message[i+6] << 48) | ((uint64_t)padded_message[i+7] << 56);
		memcpy(blocks+i/8,padded_message+i,8);
	}*/
	start = omp_get_wtime();
	for(i = 0; i < padded_length/8; i+=9)
	{

		for(y = 0; y < 5; y++)
		{
			for(x = 0; x < 5; x++)
			{
				if(x+5*y < 576/64)
				{
					//printf("Doing block %.16lx\n",blocks[x+5*y+i]);
					S[x*STATE_ROWS+y] ^= blocks[x+5*y+i];
				}
			}
		}
		for(j = 0; j < 24; j++)
	{
		
	
	
	//θ step
	for(x = 0; x < STATE_ROWS; x++)
	{
		C[x] = S[x*STATE_ROWS] ^ S[x*STATE_ROWS+1] ^ S[x*STATE_ROWS+2] ^ S[x*STATE_ROWS+3] ^ S[x*STATE_ROWS+4];
	}
	for(x = 0; x < STATE_ROWS; x++)
	{
		D[x] = C[(x+4)%STATE_ROWS] ^ rot(C[(x+1)%STATE_ROWS],1);
	}
	for(x = 0; x < STATE_ROWS; x++)
	{
		for(y = 0; y < STATE_ROWS; y++)
		{
			S[x*STATE_ROWS+y] = S[x*STATE_ROWS+y] ^ D[x];
		}
	}

	//ρ and π steps
	for(x = 0; x < 5; x++)
	{
		for(y = 0; y < 5; y++)
		{
			B[y*STATE_ROWS+((2*x+3*y)%STATE_ROWS)] = rot(S[x*STATE_ROWS+y], rot_const[x][y]);
		}
	}

	//χ step
	//temp_time = omp_get_wtime();
	
	for(x = 0; x < 5; x++)
	{
		for(y = 0; y < 5; y++)
		{
			S[x*STATE_ROWS+y] = B[x*STATE_ROWS+y] ^ ((~B[((x+1)%STATE_ROWS)*STATE_ROWS+y]) & B[((x+2)%STATE_ROWS)*STATE_ROWS+y]);
		}
	}
	//profile_time += omp_get_wtime() - temp_time;

	//ι step
	S[0] = S[0] ^ RC[j];
	}
	}
	
	count = 0;
	for(y = 0; y < 5 && count < 8; y++)
	{
		for(x = 0; x < 5 && count < 8; x++)
		{
			if(x+5*y < 576/64)
			{
				sprintf(temp,"%.16lX",reverse(S[x*STATE_ROWS+y]));
				strcat(output,temp);
				count++;
			}
		}
	}
	//keccak_func(S);
	start = omp_get_wtime() - start;
	/*for(i = 0; i < padded_length/8; i++)
	{
		printf("%.16lX\n",blocks[i]);
	}*/
	//printf("\n");
	printf("%s %s\n",output,argv[1]);
	
	printf("Time: %lf\nProfiled part: %lf\n",start,profile_time);
	return 0;
}
